"""
1.	利用Pytorch或tf2.0深度学习框架，参考yolov3的主干网络darknet53模型（参见下图），进行cifar10数据集的模型训练和模型推理，
按下面要求完成相应代码。（40分）
【For Tensorflow 2.x】
"""
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, activations, optimizers, losses, metrics, callbacks
import numpy as np
from sklearn.model_selection import train_test_split
import os
import sys

VER = 'v1.1'
ALPHA = 1e-4
BATCH_SIZE = 4
N_EPOCHS = 2
FACTOR = 20

tf.random.set_seed(1)
np.random.seed(1)
FILE_NAME = os.path.basename(__file__)
LOG_DIR = os.path.join('_log', FILE_NAME, VER)
SAVE_DIR = os.path.join('_save', FILE_NAME, VER)
SAVE_PATH = os.path.join(SAVE_DIR, 'weight')

# ①	读入cifar10数据集，并进行必要的预处理
(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()


def select_data(x, y, FACTOR):
    x_train_len = len(x)
    x_train_perm = np.random.permutation(x_train_len)
    select = x_train_perm % FACTOR == 0
    x = x[select]
    y = y[select]
    return x, y


x_train, y_train = select_data(x_train, y_train, FACTOR)
x_test, y_test = select_data(x_test, y_test, FACTOR)

x_train = x_train.astype(np.float32) / 255. * 2. - 1.
x_test = x_test.astype(np.float32) / 255. * 2. - 1.
x_test, x_val, y_test, y_val = train_test_split(x_test, y_test, train_size=0.5, random_state=1, shuffle=True)
print('x_train', x_train.shape)
print('y_train', y_train.shape)
print('x_test', x_test.shape)
print('y_test', y_test.shape)
print('x_val', x_val.shape)
print('y_val', y_val.shape)


# ②	构建卷积单元类封装ConvCell，由卷积、BN、LeakyRelu激活组成
def ConvCell(filters, ksize=(3, 3), strides=(1, 1), padding='same'):
    return keras.Sequential([
        layers.Conv2D(filters, ksize, strides, padding, use_bias=False),
        layers.BatchNormalization(),
        layers.ReLU()
    ])


# ③	构建残差块类封装ResnetBlock
class ResnetBlock(keras.Model):

    def __init__(self, n_units, out_ch, **kwargs):
        super().__init__(**kwargs)
        self.list = []
        for _ in range(n_units):
            unit = keras.Sequential([
                ConvCell(out_ch // 2, (1, 1)),
                ConvCell(out_ch, (3, 3)),
            ])
            self.list.append(unit)

    def call(self, x, training=None):
        for unit in self.list:
            x = unit(x, training=training) + x
        return x


# ④	构建yolov3的主干网络类封装DarkNet53（每个卷积模块只1个卷积单元和1个残差块组成）
class DarkNet53(keras.Model):

    def __init__(self, n_cls, **kwargs):
        super().__init__(**kwargs)
        self.conv1 = ConvCell(32)
        self.conv2 = ConvCell(64, (3, 3), (2, 2))
        self.res3_4 = ResnetBlock(1, 64)

        self.conv5 = ConvCell(128, (3, 3), (2, 2))
        self.res6_9 = ResnetBlock(2, 128)

        self.conv10 = ConvCell(256, (3, 3), (2, 2))
        self.res11_26 = ResnetBlock(8, 256)

        self.conv27 = ConvCell(512, (3, 3), (2, 2))
        self.res28_43 = ResnetBlock(8, 512)

        self.conv44 = ConvCell(1024, (3, 3), (2, 2))
        self.res45_52 = ResnetBlock(4, 1024)

        self.avg = layers.GlobalAveragePooling2D()

        self.fc53 = layers.Dense(n_cls)

    def call(self, x, training=None):
        x = self.conv1(x, training=training)
        x = self.conv2(x, training=training)
        x = self.res3_4(x, training=training)
        x = self.conv5(x, training=training)
        x = self.res6_9(x, training=training)
        x = self.conv10(x, training=training)
        x = self.res11_26(x, training=training)
        x = self.conv27(x, training=training)
        x = self.res28_43(x, training=training)
        x = self.conv44(x, training=training)
        x = self.res45_52(x, training=training)
        x = self.avg(x, training=training)
        x = self.fc53(x, training=training)
        return x


model = DarkNet53(10)
model.build((None, 32, 32, 3))

# ⑥	打印输出整体模型维度结构
model.summary()

# ⑤	进行前向传播
model.compile(
    loss=losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=optimizers.Adam(learning_rate=ALPHA),
    metrics=[
        metrics.SparseCategoricalAccuracy(),
    ],
)

# ⑦	训练集数据进行模型训练
if os.path.exists(SAVE_DIR):
    print('Loading...')
    model.load_weights(SAVE_PATH)
    print('Loaded')
else:
    model.fit(x_train,
              y_train,
              BATCH_SIZE,
              N_EPOCHS,
              validation_data=(x_val, y_val),
              # validation_batch_size=BATCH_SIZE,
              callbacks=[callbacks.TensorBoard(log_dir=LOG_DIR, update_freq='batch', profile_batch=0)],
              )
    print('Saving...')
    os.makedirs(SAVE_DIR, exist_ok=True)
    model.save_weights(SAVE_PATH)
    print('Saved')

# ⑧	模型训练完毕后，对比打印输出训练集和测试集的准确率
print('Testing ...')
model.evaluate(x_test,
               y_test,
               BATCH_SIZE)
print('Tested')
